# ------------------------------------------------------------------------------
# Color Tracking Detector
# ------------------------------------------------------------------------------
# 基于颜色的物体追踪
# ------------------------------------------------------------------------------

import cv2
import numpy as np
from .base_detector import BaseDetector


class ColorTracker(BaseDetector):
    """颜色追踪检测器"""
    
    def __init__(self, hsv_min=None, hsv_max=None):
        super().__init__()
        # 默认追踪绿色物体
        self.hsv_min = hsv_min if hsv_min is not None else np.array([50, 80, 80])
        self.hsv_max = hsv_max if hsv_max is not None else np.array([120, 255, 255])
        
    def update_color_range(self, hsv_min, hsv_max):
        """更新颜色范围"""
        self.hsv_min = np.array(hsv_min)
        self.hsv_max = np.array(hsv_max)
        
    def process_frame(self, frame):
        """
        追踪指定颜色的物体
        
        Args:
            frame: 输入图像帧
            
        Returns:
            display_frame: 显示帧（原图 + 遮罩）
        """
        self.update_fps()
        
        # 应用模糊
        blurred = cv2.blur(frame, (3, 3))
        
        # 转换到 HSV 色彩空间
        hsv = cv2.cvtColor(blurred, cv2.COLOR_BGR2HSV)
        
        # 创建颜色遮罩
        mask = cv2.inRange(hsv, self.hsv_min, self.hsv_max)
        
        # 查找轮廓
        contours, _ = cv2.findContours(mask, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
        
        # 找到最大的轮廓
        max_area = 0
        best_cnt = None
        for cnt in contours:
            area = cv2.contourArea(cnt)
            if area > max_area:
                max_area = area
                best_cnt = cnt
        
        # 绘制中心点和边界
        if best_cnt is not None and max_area > 100:  # 最小面积阈值
            M = cv2.moments(best_cnt)
            if M['m00'] != 0:
                cx = int(M['m10'] / M['m00'])
                cy = int(M['m01'] / M['m00'])
                
                # 绘制中心点
                cv2.circle(frame, (cx, cy), 5, (0, 0, 255), -1)
                
                # 绘制轮廓
                cv2.drawContours(frame, [best_cnt], -1, (0, 255, 0), 2)
                
                # 显示坐标
                coord_text = f'中心: ({cx}, {cy})'
                cv2.putText(frame, coord_text, (cx + 10, cy), 
                           cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
                
                # 显示面积
                area_text = f'面积: {int(max_area)}'
                cv2.putText(frame, area_text, (24, 50), 
                           cv2.FONT_HERSHEY_PLAIN, 1, (0, 255, 0), 1)
        else:
            # 未找到目标
            cv2.putText(frame, '未检测到目标', (24, 50), 
                       cv2.FONT_HERSHEY_PLAIN, 1, (0, 0, 255), 1)
        
        # 创建显示帧（原图 + 遮罩）
        mask_colored = cv2.cvtColor(mask, cv2.COLOR_GRAY2BGR)
        display_frame = cv2.hconcat([frame, mask_colored])
        
        # 显示当前 HSV 范围
        hsv_text = f'HSV: [{self.hsv_min[0]}-{self.hsv_max[0]}, {self.hsv_min[1]}-{self.hsv_max[1]}, {self.hsv_min[2]}-{self.hsv_max[2]}]'
        cv2.putText(display_frame, hsv_text, (24, 70), 
                   cv2.FONT_HERSHEY_PLAIN, 0.8, (255, 255, 0), 1)
        
        display_frame = self.visualize_fps(display_frame)
        return display_frame

